package Tree;

import java.util.Arrays;

public class KthAncestor {
    static class TreeAncestor {
        public static int MAXN = 50001;
        public static int LIMIT = 16;

        public static int power;

        public static int log2(int n) {
            int ans = 0;
            while ((1 << ans) <= (n >> 1)) {
                ans++;
            }
            return ans;
        }

        public static int[] head = new int[MAXN];

        public static int[] next = new int[MAXN];

        public static int[] to = new int[MAXN];

        public static int cnt;

        public static int[] deep = new int[MAXN];

        public static int[][] stjump = new int[MAXN][LIMIT];

        public TreeAncestor(int n, int[] parent) {
            power = log2(n);
            cnt = 1;
            Arrays.fill(head, 0, n, 0);
            for (int i = 1; i < parent.length; i++) {
                addEdge(parent[i], i);
            }
            // 从根节点开，搜集每个节点的深度信息和向上跳2^n到达的点信息
            dfs(0, 0);
        }

        private void addEdge(int u, int v) {
            // 记录当前边的上一条边号
            next[cnt] = head[u];
            // 记录当前边去往的节点
            to[cnt] = v;
            // 记录新的边号
            head[u] = cnt++;
        }

        public static void dfs(int i, int f) {
            if (i == 0) {
                // 根节点深度为1
                deep[i] = 1;
            } else {
                // 当前节点的深度 = 父节点深度 + 1
                deep[i] = deep[f] + 1;
            }
            // 向上跳2^0层 到达(=) 父节点(f)
            stjump[i][0] = f;
            // 跳出条件：向上跳的步数 > 当前节点的深度
            // 所以: 根节点一步都不能跳
            for (int p = 1; (1 << p) <= deep[i]; p++) {
                // 向上跳2^p层 到达(=) 向上跳2^(p-1)层后，在向上跳2^(p-1)层到达的节点
                stjump[i][p] = stjump[stjump[i][p - 1]][p - 1];
            }
            for (int e = head[i]; e != 0; e = next[e]) {
                // 递归处理 i 节点的子节点
                dfs(to[e], i);
            }
        }

        /**
         * @param i 树节点
         * @param k 向上走k步
         * @return 到达的节点
         */
        public int getKthAncestor(int i, int k) {
            if (deep[i] <= k) {
                // 族谱没有这么长，找不到
                return -1;
            }

            // s 是想要去往的层数 = i 所在的层数 - 向上走的步数
            int s = deep[i] - k;

            for (int p = power; p >= 0; p--) {
                if (deep[stjump[i][p]] >= s) {
                    // 向上条2^p层，还是在目标层(s)的下面，就先跳过去
                    // 在向上接着跳
                    i = stjump[i][p];
                }
            }
            return i;
        }
    }
}
